import joblib
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.model_selection import GridSearchCV
from sklearn.metrics import mean_squared_error, mean_absolute_error, root_mean_squared_error, accuracy_score, \
    roc_auc_score, f1_score
from sklearn.preprocessing import StandardScaler


def train_model():
    data = pd.read_csv("../../data/processed/train_final.csv")
    x = data.iloc[:, :-1]
    y = data.iloc[:, -1]

    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.2, random_state=18)

    scaler = StandardScaler()
    x_train_scaled = scaler.fit_transform(x_train)  # 对平衡后的训练集标准化
    x_test_scaled = scaler.transform(x_test)

    es = RandomForestClassifier(n_estimators=100, random_state=18)
    # 1. 定义要搜索的参数网格
    param_grid = {
        'n_estimators': [50, 100, 200],  # 树的数量
        'max_depth': [10, 20, 30],
        'min_samples_split': [2, 5, 10],  # 分裂所需最小样本数
        'min_samples_leaf': [1, 2, 4]  # 叶节点最小样本数
    }

    # 2. 初始化GridSearchCV对象
    # estimator: 基础模型
    # param_grid: 参数网格
    # cv: 5折交叉验证
    # scoring: 评估指标，分类常用'accuracy'，回归常用'r2'
    # n_jobs: 并行数，-1表示使用所有CPU核心
    grid_search = GridSearchCV(
        estimator=es,
        param_grid=param_grid,
        cv=5,
        scoring='roc_auc',
    )

    # 3. 在训练数据上执行网格搜索（这个过程可能很耗时）
    grid_search.fit(x_train_scaled, y_train)

    # 4. 输出最佳参数和最佳得分
    print("最佳参数组合: ", grid_search.best_params_)
    print("最佳交叉验证得分: ", grid_search.best_score_)

    # 5. 直接使用最佳参数的模型进行预测
    best_rf_model = grid_search.best_estimator_
    y_pred = best_rf_model.predict(x_test_scaled)
    y_pred_proba = grid_search.predict_proba(x_test_scaled)[:, 1]
    # print("调优后测试集准确率: ", accuracy_score(y_test, y_pred_best))
    print(f"auc值：{roc_auc_score(y_test, y_pred_proba)}")
    print(accuracy_score(y_pred, y_test))
    print(f"F1值：{f1_score(y_test, y_pred)}")
    joblib.dump(best_rf_model, "../../model/人才流失.pkl")


# def model_todo(path):


if __name__ == '__main__':
    train_model()
